import json
import argparse
import regex
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument("--files", type=str, help='parse file list')
parser.add_argument("--to_files", type=str, help='write file list')
args = parser.parse_args()


def parse_json_log(files):
    res = []
    pat = r'\[.*?\]'
    for f in files:
        iters = 0
        lengths = 0
        with open(f, "r", encoding="utf-8")as  fp:
            mp = json.load(fp)
            lt = []
            res.append(lt)
            for item in tqdm(mp):
                if 'ph' in item and item['ph'] == 'E' and 'tid' in item and item['tid'].strip().startswith('gpu0:1') and \
                        item['name'].strip() != 'BatchNorm':
                    tmp = regex.findall(pat, item['args']['__outputs__'])[0][1:-1].split(",")
                    val = 1
                    for t in tmp:
                        val *= int(t.strip())
                    lt.append((item['ts'], val))
                    lengths += val
                elif item['ph'] == 'E' and item['tid'].strip().startswith('gpu0:2'):
                    lt.append(('-'*10, '{} iter_end: {} total length: {}'.format(item['ts'], iters, lengths)))
                    iters += 1
                    lengths = 0
    return res


if __name__ == '__main__':
    files = args.files.strip().split(",")
    files = [f.strip() for f in files]
    res = parse_json_log(files)
    to_files = args.to_files.strip().split(",")
    to_files = [f.strip() for f in to_files]
    for i, f in enumerate(to_files):
        with open(f, "w", encoding="utf-8")as fp:
            for item in res[i]:
                fp.write(f"{item[0]} {item[1]}\n")
